#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <mpi.h>

#define MIN_SIZE 128
#define MAX_SIZE 2048

// 生成随机矩阵
void generate_matrix(double *matrix, int rows, int cols) {
    for (int i = 0; i < rows * cols; i++) {
        matrix[i] = (double)rand() / RAND_MAX * 10.0;
    }
}

// 打印矩阵（用于调试）
void print_matrix(double *matrix, int rows, int cols) {
    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < cols; j++) {
            printf("%.2f ", matrix[i * cols + j]);
        }
        printf("\n");
    }
}

// 矩阵乘法（局部计算）
void matrix_multiply(double *A, double *B, double *C, 
                    int m, int n, int k, 
                    int start_row, int end_row) {
    for (int i = start_row; i < end_row; i++) {
        for (int j = 0; j < k; j++) {
            C[i * k + j] = 0.0;
            for (int l = 0; l < n; l++) {
                C[i * k + j] += A[i * n + l] * B[l * k + j];
            }
        }
    }
}

int main(int argc, char *argv[]) {
    int rank, size;
    int m, n, k;
    double *A = NULL, *B = NULL, *C = NULL;
    double start_time, end_time;
    
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    
    // 设置随机种子（主进程）
    if (rank == 0) {
        srand(time(NULL));
    }
    
    // 矩阵尺寸
    if (rank == 0) {
    
        m = 512; n = 512; k = 512;
        
        // 验证矩阵尺寸
        if (m < MIN_SIZE || m > MAX_SIZE || 
            n < MIN_SIZE || n > MAX_SIZE || 
            k < MIN_SIZE || k > MAX_SIZE) {
            printf("Matrix size out of range [128, 2048]\n");
            MPI_Abort(MPI_COMM_WORLD, 1);
        }
    }
    
    // 广播矩阵尺寸
    MPI_Bcast(&m, 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&n, 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&k, 1, MPI_INT, 0, MPI_COMM_WORLD);
    
    // 主进程生成矩阵A和B
    if (rank == 0) {
        A = (double *)malloc(m * n * sizeof(double));
        B = (double *)malloc(n * k * sizeof(double));
        C = (double *)malloc(m * k * sizeof(double));
        
        generate_matrix(A, m, n);
        generate_matrix(B, n, k);
        
        // 打印部分矩阵（调试用）
        printf("Matrix A (partial):\n");
        print_matrix(A, 2, 2);
        printf("\nMatrix B (partial):\n");
        print_matrix(B, 2, 2);
    } else {
        // 其他进程分配B矩阵的空间
        B = (double *)malloc(n * k * sizeof(double));
    }
    
    // 广播B矩阵到所有进程
    MPI_Bcast(B, n * k, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    
    // 计算每个进程负责的行范围
    int rows_per_proc = m / size;
    int remainder = m % size;
    int start_row = rank * rows_per_proc;
    int end_row = start_row + rows_per_proc;
    
    // 处理不能整除的情况
    if (rank < remainder) {
        start_row += rank;
        end_row += rank + 1;
    } else {
        start_row += remainder;
        end_row += remainder;
    }
    
    // 分配本地A和C的空间
    double *local_A = (double *)malloc((end_row - start_row) * n * sizeof(double));
    double *local_C = (double *)malloc((end_row - start_row) * k * sizeof(double));
    
    // 主进程分发A的行
    if (rank == 0) {
        // 首先复制自己的部分
        for (int i = 0; i < end_row - start_row; i++) {
            for (int j = 0; j < n; j++) {
                local_A[i * n + j] = A[(start_row + i) * n + j];
            }
        }
        
        // 发送其他部分给其他进程
        for (int dest = 1; dest < size; dest++) {
            int dest_start = dest * rows_per_proc;
            int dest_end = dest_start + rows_per_proc;
            
            if (dest < remainder) {
                dest_start += dest;
                dest_end += dest + 1;
            } else {
                dest_start += remainder;
                dest_end += remainder;
            }
            
            MPI_Send(&A[dest_start * n], (dest_end - dest_start) * n, 
                    MPI_DOUBLE, dest, 0, MPI_COMM_WORLD);
        }
    } else {
        // 接收来自主进程的A的行
        MPI_Recv(local_A, (end_row - start_row) * n, 
                MPI_DOUBLE, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }
    
    // 开始计时
    MPI_Barrier(MPI_COMM_WORLD);
    start_time = MPI_Wtime();
    
    // 本地矩阵乘法计算
    matrix_multiply(local_A, B, local_C, 
                   m, n, k, 
                   0, end_row - start_row);
    
    // 结束计时
    end_time = MPI_Wtime();
    
    // 收集结果到主进程
    if (rank == 0) {
        // 复制自己的部分
        for (int i = 0; i < end_row - start_row; i++) {
            for (int j = 0; j < k; j++) {
                C[(start_row + i) * k + j] = local_C[i * k + j];
            }
        }
        
        // 接收其他进程的结果
        for (int src = 1; src < size; src++) {
            int src_start = src * rows_per_proc;
            int src_end = src_start + rows_per_proc;
            
            if (src < remainder) {
                src_start += src;
                src_end += src + 1;
            } else {
                src_start += remainder;
                src_end += remainder;
            }
            
            MPI_Recv(&C[src_start * k], (src_end - src_start) * k, 
                    MPI_DOUBLE, src, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        }
        
        // 打印部分结果（调试用）
        printf("\nResult Matrix C (partial):\n");
        print_matrix(C, 2, 2);
        
        // 打印时间
        printf("\nMatrix multiplication time: %.6f seconds\n", end_time - start_time);
    } else {
        // 发送结果给主进程
        MPI_Send(local_C, (end_row - start_row) * k, 
                MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
    }
    
    // 释放内存
    free(local_A);
    free(local_C);
    if (rank == 0) {
        free(A);
        free(B);
        free(C);
    } else {
        free(B);
    }
    
    MPI_Finalize();
    return 0;
}